from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, AIMessagePromptTemplate, \
    ChatPromptTemplate, PromptTemplate
from langchain.schema import HumanMessage, AIMessage, SystemMessage, LLMResult, ChatGeneration
from langchain.schema.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

chat = ChatOpenAI()
chat([HumanMessage(content="Translate this sentence from English to French: I love programming.")])
AIMessage(content="J'aime programmer.", additional_kwargs={})
messages = [
    SystemMessage(content="You are a helpful assistant that translates English to French."),
    HumanMessage(content="I love programming.")
]
chat(messages)
AIMessage(content="J'aime programmer.", additional_kwargs={})
batch_messages = [
    [
        SystemMessage(content="You are a helpful assistant that translates English to French."),
        HumanMessage(content="I love programming.")
    ],
    [
        SystemMessage(content="You are a helpful assistant that translates English to French."),
        HumanMessage(content="I love artificial intelligence.")
    ],
]
result = chat.generate(batch_messages)
llm_result = LLMResult(generations=[[ChatGeneration(text="J'aime programmer.", generation_info=None,
                                       message=AIMessage(content="J'aime programmer.", additional_kwargs={}))], [
                           ChatGeneration(text="J'aime l'intelligence artificielle.", generation_info=None,
                                          message=AIMessage(content="J'aime l'intelligence artificielle.",
                                                            additional_kwargs={}))]],
          llm_output={'token_usage': {'prompt_tokens': 57, 'completion_tokens': 20, 'total_tokens': 77}})

print(result.llm_output)
print(llm_result)

template = "You are a helpful assistant that translates {input_language} to {output_language}."
system_message_prompt = SystemMessagePromptTemplate.from_template(template=template)
example_human = HumanMessagePromptTemplate.from_template("你好啊，老师！")
example_ai = AIMessagePromptTemplate.from_template("Argh me mateys")
human_template = "{text}"
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)

chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, example_human, example_ai, human_message_prompt])
chain = LLMChain(llm=chat, prompt=chat_prompt)
chat(chat_prompt.format_prompt(input_language="English", output_language="French", text="I love programming.").to_messages())

prompt = PromptTemplate(
    template="You are a helpful assistant that translates {input_language} to {output_language}.",
    input_variables=["input_language", "output_language"],
)
system_message_prompt = SystemMessagePromptTemplate(prompt=prompt)


chat = ChatOpenAI(streaming=True, callbacks=[StreamingStdOutCallbackHandler()], temperature=0)
resp = chat([HumanMessage(content="Write me a song about sparkling water.")])


